#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on 5/4/2025

@author: ifenty
"""

import numpy as np
from pathlib import Path
import xarray as xr
import matplotlib.pyplot as plt;

import numpy as np
import cartopy.crs as ccrs
import cartopy.feature as cfeature

def calc_adxx_length(adxx):
    tmp = []
    # number of time levels
    for i in range(adxx.shape[0]):
        tmpv = np.inner(np.ravel(adxx[i]), np.ravel(adxx[i]))
        tmp.append(tmpv)

    return np.array(tmp)

def plot_adxx_length(adxx, lags, t=''):
    plt.figure()
    adxx_len = calc_adxx_length(adxx)
    plt.plot(lags,adxx_len,'k.-')
    plt.title('|adxx ' + t + '| vs. lag')

def plot_gradient_field(adxx_DA, central_longitude=-180):
    
    field_to_plot_max = np.nanmax(np.abs(np.ravel(adxx_DA)))*0.6
    
    # Plot
    fig = plt.figure(figsize=(15, 5))
    ax = plt.axes(projection=ccrs.PlateCarree(central_longitude=central_longitude))
    
    pcm = ax.pcolormesh(adxx_DA.lon, adxx_DA.lat, adxx_DA, 
                        vmin = -field_to_plot_max, 
                        vmax = field_to_plot_max, 
                        transform=ccrs.PlateCarree(), cmap="RdBu_r")
    
    ax.coastlines()
    ax.gridlines(draw_labels=True)
    plt.colorbar(pcm, orientation='vertical')
    plt.title(adxx_DA.name + ' at lag ' + str(adxx_DA.lag.values))
    plt.show()

Define Paths

Paths of EMU gradient experiments

emu_output_basedir= Path('/home/ifenty/incoming/gradients')
gradient_dirs = ['emu_adj_48_48_1_15_743_1','emu_adj_48_48_1_2_209_1','emu_adj_48_48_1_56_1069_1','emu_adj_48_48_1_56_1072_1']

Find all EMU adjoint experiment directories

gd_output_dirs = []
gd_files = dict()
for gd in gradient_dirs:
    tmp = emu_output_basedir / gd / 'output'
    print(tmp)
    gd_output_dirs.append(tmp)
    gd_files[gd] = np.sort(list(Path(tmp).glob('*nc')))
    print(gd_files[gd])

adxx_files = ['adxx_qnet','adxx_tauu','adxx_tauv']

/home/ifenty/incoming/gradients/emu_adj_48_48_1_15_743_1/output
[PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_15_743_1/output/adxx_qnet.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_15_743_1/output/adxx_tauu.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_15_743_1/output/adxx_tauv.nc')]
/home/ifenty/incoming/gradients/emu_adj_48_48_1_2_209_1/output
[PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_2_209_1/output/adxx_qnet.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_2_209_1/output/adxx_tauu.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_2_209_1/output/adxx_tauv.nc')]
/home/ifenty/incoming/gradients/emu_adj_48_48_1_56_1069_1/output
[PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_56_1069_1/output/adxx_qnet.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_56_1069_1/output/adxx_tauu.nc')
 PosixPath('/home/ifenty/incoming/gradients/emu_adj_48_48_1_56_1069_1/output/adxx_tauv.nc')]
/home/ifenty/incoming/gradients/emu_adj_48_48_1_56_1072_1/output
[]

# loop through each EMU adjoint gradient experiment directory
# directory 3 and 4 are almost identical, so skip 4

# choose an EMU experiment to process
for dd in range(3):

    # for plotting, select the central longitude
    if dd == 0: # tropical pacific experiment
        central_longitude = -180 # degrees east
    else: # Nantucket and N. Atlantic experiments
        central_longitude = -45 # degrees east

    # make maps of the gradients for several lag
    lags_to_show = [-1, -6, -36, -72]

    adxx_qnet_DA = xr.open_dataset(gd_output_dirs[dd] / 'adxx_qnet.nc')['adxx_qnet']
    adxx_tauu_DA = xr.open_dataset(gd_output_dirs[dd] / 'adxx_tauu.nc')['adxx_tauu']
    adxx_tauv_DA = xr.open_dataset(gd_output_dirs[dd] / 'adxx_tauv.nc')['adxx_tauv']    

    # plot gradients for qnet, tauu, tauv
    for adxx_DA in [adxx_tauu_DA, adxx_tauv_DA, adxx_qnet_DA]:
        
        # plot different lags
        for ll in lags_to_show:
            plot_gradient_field(adxx_DA[ll], central_longitude=central_longitude)

        # plot the length of the gradient vector (crudely calculated)
        tmp = calc_adxx_length(np.where(np.isnan(adxx_DA.values),0, adxx_DA.values))
        plot_adxx_length(tmp[-72:], adxx_DA.lag[-72:], adxx_DA.name);plt.grid()
        plt.xlabel('days')
        plt.ylabel('gradient length')

    # close the DataArray file
    adxx_qnet_DA.close()
    adxx_tauu_DA.close()
    adxx_tauv_DA.close()